package middleware

import (
	"context"
	"github.com/zeromicro/go-zero/core/logx"
	"mall/common/excption"
	"mall/common/limit"
	"mall/common/response"
	"net/http"
)

type LimitRateMiddleware struct {
	limit        *limit.PeriodLimit
	rateLimitKey string
}

func NewLimitRateMiddleware(periodLimit *limit.PeriodLimit, rateLimitKey string) *LimitRateMiddleware {
	return &LimitRateMiddleware{limit: periodLimit, rateLimitKey: rateLimitKey}
}

func (l *LimitRateMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(writer http.ResponseWriter, request *http.Request) {
		takeCode, err := l.limit.TakeCtx(context.Background(), l.rateLimitKey)
		if err != nil {
			response.Response(writer, nil, err)
			return
		}
		limitErr := doRateLimit(takeCode, l.rateLimitKey)
		if limitErr != nil {
			response.Response(writer, nil, limitErr)
			return
		}
		// pass
		next(writer, request)
	}
}

func doRateLimit(takeCode int, rateLimitKey string) error {
	switch takeCode {
	case limit.OverQuota:
		logx.Errorf("OverQuota key: %v", rateLimitKey)
		return excption.NewErrCode(excption.RequestLimitError)
	case limit.Allowed:
		fallthrough
	case limit.HitQuota:
		logx.Infof("AllowedQuota key: %v", rateLimitKey)
		return nil
	default:
		logx.Errorf("DefaultQuota key: %v", rateLimitKey)
		// unknown response, we just let the sms go
		return excption.NewErrCode(excption.RequestLimitError)
	}
}
